package com.atguigu.userprofile.ml.pipline

import org.apache.spark.ml.classification.{DecisionTreeClassificationModel, DecisionTreeClassifier}
import org.apache.spark.ml.feature.{StringIndexer, VectorAssembler, VectorIndexer}
import org.apache.spark.ml.{Pipeline, PipelineModel, Transformer}
import org.apache.spark.sql.DataFrame

class MyPipeline {


    var pipeline:Pipeline = null;

    def init():MyPipeline={
        pipeline = new Pipeline().setStages(Array(
            createLableIndexer(),
            createFeatureAssemble(),
            createFeatureIndexer(),
            createClassfier()
        ))
        this
    }

    var pipelineModel:PipelineModel = null;

    var labelColName:String = null

    def setLableColName(labelColName:String):MyPipeline={
        this.labelColName = labelColName
        this
    }


    var featureColName:Array[String] = null

    def setFeatureColNames(featureColName:Array[String]):MyPipeline ={
        this.featureColName=featureColName
        this
    }

    var maxCategories:Int = 10
    def setMaxCategories(maxCategories:Int):MyPipeline ={
        this.maxCategories=maxCategories
        this
    }

    // 1)标签索引：
    // 专门整理标签数据（参考答案）
    //     1.提取标签值
    //     2.把标签值按照出现的比例，变为矢量值（0,1,2,3.。。） 出现的比例越高，矢量值越小
    // 参数 inputcol :原始标签列名称
    //     outputcol:新产生矢量列的名称
    def createLableIndexer():StringIndexer ={
        val indexer = new StringIndexer()
        indexer.setInputCol(labelColName)
        indexer.setOutputCol("label_index")

        indexer
    }

    // 2)特征聚合
    //   1.提取特征
    //   2.把多个特征列组合成为一个特征列
    // 参数 inputcol :原始特征列名称（多个）
    //     outputcol:聚合特征列

    def createFeatureAssemble():VectorAssembler ={
        val vectorAssembler = new VectorAssembler()
        vectorAssembler.setInputCols(featureColName)
        vectorAssembler.setOutputCol("feature_assemble")
        vectorAssembler
    }

    // 3)特征索引
    //  把特征  向量化处理
    //  1.把原始值转为向量值(0,1,2,3...) 出现的比例越高，矢量值越小
    //  连续值特征不会被向量化处理 维持原值
    // 参数 inputcol :上一步的输出列
    //     outputcol:新建的特征向量列
    def createFeatureIndexer():VectorIndexer={
        val vecrorIndexer = new VectorIndexer()
        vecrorIndexer.setInputCol("feature_assemble")
        vecrorIndexer.setOutputCol("feature_index")
        vecrorIndexer.setMaxCategories(maxCategories)  //区分数据是否为连续值
        vecrorIndexer
    }


    // 4)分类器
    //  决策树算法
    def createClassfier():DecisionTreeClassifier={
        val classifier = new DecisionTreeClassifier()
        classifier.setLabelCol("label_index")           //标签向量列
        classifier.setFeaturesCol("feature_index")      //特征向量列
        classifier.setPredictionCol("prediction")       //预测结果列
        classifier.setImpurity("gini")                  //默认为gini
        // classifier参数可以优化
        classifier
    }


    //训练
    //传入数据  （含标签） + 流水线 ==> 流水线模型
    def train(dataFrame: DataFrame)={
       pipelineModel = pipeline.fit(dataFrame)
    }

    //预测
    def predict(dataFrame: DataFrame):DataFrame ={
        val predicteddataFrame: DataFrame = pipelineModel.transform(dataFrame)
        predicteddataFrame
    }

    //打印决策树
    def printTree()={
        val transformer: Transformer = pipelineModel.stages(3)
        val classificationModel: DecisionTreeClassificationModel = transformer.asInstanceOf[DecisionTreeClassificationModel]
        println(classificationModel.toDebugString)

    }

    //为了更好地筛选特征，打印特征的权重
    def printFeatureWeight()={
        val transformer: Transformer = pipelineModel.stages(3)
        val classificationModel: DecisionTreeClassificationModel = transformer.asInstanceOf[DecisionTreeClassificationModel]
        println(classificationModel.featureImportances)
    }
}
